import os
import json
from dotenv import load_dotenv

from .xfyun_chat import XFYunChat

# 加载 .env 环境变量
load_dotenv(os.path.join(os.path.dirname(os.getcwd()), ".env"))
load_dotenv()

# 初始化讯飞大模型
kg_model = XFYunChat(
    appid=os.getenv("XFYUN_APPID"),
    api_key=os.getenv("XFYUN_API_KEY"),
    api_secret=os.getenv("XFYUN_API_SECRET"),
    system_prompt="你是一个专业的知识图谱专家，能够围绕任何学科知识点生成高质量、结构化的三元组。"
)

# 🌱 单次生成函数
def generate_knowledge_triples(entity: str):
    prompt = f"""
请围绕知识点“{entity}”生成适合教学和学习使用的高质量知识图谱三元组（每行一个标准 JSON）：
{{"h": "当前知识点", "r": "关系类型", "t": "相关知识点"}}

要求：
1. 请先判断“{entity}”所属的学科领域（如数学、物理、化学、文学、历史、地理、生物、计算机、音乐等），并根据该领域知识体系生成三元组。
2. 每行输出一个标准 JSON 格式三元组，每个三元组都与“{entity}”直接相关。
3. 每个三元组字段要求如下：
   - "h"：必须是简洁、准确的知识点名称，不得是完整句子或带解释性内容。
   - "r"：必须是一个单词或短语，只能选用以下关键词之一：
     ["定义", "性质", "单位", "公式", "应用", "实验", "历史", "相关概念", "常见问题", "例题", "组成部分", "类别", "类型", "对比", "类比"]
   - "t"：必须是简洁、准确的相关知识点名称，不得是完整句子或带解释性内容。
4. 每个关系（"r"）只写一个关键词，不允许用 "/" 连接多个。
5. 至少生成 10 条高质量三元组，覆盖不同的关系类型。
6. 所有输出必须是准确、权威、适合教学使用的结构化内容。
7. 输出必须是纯粹的标准 JSON 行，不要任何多余解释说明或格式错误。
"""
    response = kg_model.chat(prompt)
    triples = []
    for line in response.strip().split('\n'):
        try:
            triples.append(json.loads(line.strip()))
        except json.JSONDecodeError:
            continue
    return triples


# 🌱 主构图逻辑：递归扩展知识图谱
ALLOWED_RELATIONS = {
    "相关概念", "上位概念", "下位概念", "组成部分", "类别", "是", "属于", "类型", "对比", "类比", "例题", "常见问题"
}

def build_knowledge_graph(seed_entity: str, max_depth: int = 2, max_nodes: int = 50):
    visited = set()
    queue = [(seed_entity, 0)]
    all_triples = []

    while queue and len(visited) < max_nodes:
        current_entity, depth = queue.pop(0)
        if current_entity in visited or depth > max_depth:
            continue

        print(f" 正在扩展：{current_entity} (深度：{depth})")
        triples = generate_knowledge_triples(current_entity)
        visited.add(current_entity)

        for triple in triples:
            all_triples.append(triple)
            next_entity = triple.get("t")
            relation = triple.get("r")

            #  只扩展允许的关系
            if relation in ALLOWED_RELATIONS and next_entity and next_entity not in visited:
                queue.append((next_entity, depth + 1))

    return all_triples


# 🌟 运行入口
if __name__ == "__main__":
    seed = input(" 请输入知识起始词：").strip()
    graph = build_knowledge_graph(seed, max_depth=2, max_nodes=50)

    # 保存
    output_path = f"kg_{seed}_expanded.jsonl"
    with open(output_path, "w", encoding="utf-8") as f:
        for triple in graph:
            f.write(json.dumps(triple, ensure_ascii=False) + "\n")

    print(f"\n✅ 知识图谱扩展完成，共生成 {len(graph)} 条三元组")
    print(f"📁 输出已保存到：{output_path}")
